import os
import torch
from modelscope import snapshot_download
from diffsynth import ModelManager, FluxImagePipeline

# 基础配置
MODEL_NAME = 'yiwanji/FLUX_xiao_hong_shu_ji_zhi_zhen_shi_V2'
MODEL_DIR = './models'  # 本地模型存储路径
OUTPUT_DIR = './output'  # 输出目录
PROMPT = "a watercolor painting of a cat"  # 简化提示词


def main():
    # 自动选择设备（优先CPU以降低显存需求）
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    # 1. 模型下载
    model_path = snapshot_download(MODEL_NAME,
                                   cache_dir=MODEL_DIR,
                                   revision='v1.0')

    # 2. 初始化模型管理器（使用半精度节省内存）
    manager = ModelManager(torch_dtype=torch.float16)

    # 3. 仅加载必需组件
    manager.load_models(
        file_path_list=[model_path],
        model_names=['text_encoder', 'unet', 'vae'],  # 最小组件集
        device=device
    )

    # 4. 创建低内存优化pipeline
    pipe = FluxImagePipeline.from_model_manager(
        manager,
        device=device,
        torch_dtype=torch.float16
    ).to(device)

    # 5. 内存优化配置
    pipe.enable_cpu_offload()  # CPU卸载
    # if hasattr(pipe, 'dit'):
        # pipe.dit.quantize()  # 模型量化

    # 6. 生成图像（简化参数）
    image = pipe(
        prompt=PROMPT,
        num_inference_steps=10,  # 减少步数以加快生成
        seed=42,  # 固定随机种子保证可重复性
        # scale=5.0  # 更保守的指导系数
    )

    # 保存结果
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    image.save(f"{OUTPUT_DIR}/output.jpg")
    print("Image generated successfully!")


if __name__ == "__main__":
    main()